import numpy as np
from numba import njit
import pandas as pd 
import pybroker
from pybroker import Strategy
import talib  
import matplotlib.pyplot as plt
from pybroker.ext.data import AKShare
pybroker.enable_data_source_cache('AKShare')

def ma(bar_data,short_window,long_window):
    #@njit  # Enable Numba JIT.
    def vec_ma(values):
        # Initialize the result array.
        n = len(values)
        
        data = np.full((n, 2), np.nan)  
        # 将NumPy数组转换为pandas DataFrame，并添加列标签  
        df = pd.DataFrame(data, columns=['SMA_short', 'SMA_long'])  

        # 计算短期（例如5日）和长期（例如20日）均线  
        df['SMA_short'] = talib.SMA(values, timeperiod=short_window)  
        df['SMA_long'] = talib.SMA(values, timeperiod=long_window)  
            # 初始化金叉和死叉的列  
        df['Cross'] = 0  
        out = np.array([np.nan for _ in range(n)])

        # For all bars starting at lookback:
        # 遍历DataFrame，判断金叉和死叉  
        for i in range(1, len(df)):  
            if df['SMA_short'].iloc[i] > df['SMA_long'].iloc[i] and df['SMA_short'].iloc[i-1] < df['SMA_long'].iloc[i-1]:  
                #df.at[df.index[i], 'Cross'] = 2 
                out[i] =1
            if df['SMA_short'].iloc[i] < df['SMA_long'].iloc[i] and df['SMA_short'].iloc[i-1] > df['SMA_long'].iloc[i-1]:  
                #df.at[df.index[i], 'Cross'] = 1  
                out[i] = -1
        
        print (df)
        print (out)
        return out
    
    return vec_ma(bar_data.close) # 把bar_data.close传递给values,最终out[]返回

aKShare = AKShare()
df = aKShare.query('000001', '3/1/2024', '6/7/2024')
out=ma(df,short_window=5,long_window=20)
print(df)